import redis
import mysql.connector
from contextlib import contextmanager

from utils.common import data_preprocessing


class PowerLoadModel:
    def __init__(self, path):
        # 原有初始化代码...
        self.data_source = data_preprocessing(path)
        self.time_key_load = self.data_source.set_index("time")["power_load"].to_dict()

        # Redis连接配置
        self.redis_client = redis.Redis(
            host='localhost',
            port=6379,
            db=0,
            decode_responses=True
        )

        # MySQL连接配置
        self.mysql_config = {
            'host': 'localhost',
            'user': 'root',
            'password': 'password',
            'database': 'power_load_db',
            'charset': 'utf8mb4'
        }

    @contextmanager
    def get_mysql_connection(self):
        """MySQL连接上下文管理器"""
        conn = mysql.connector.connect(**self.mysql_config)
        try:
            yield conn
        finally:
            conn.close()

    def cache_features(self, key, features):
        """将特征数据缓存到Redis"""
        import json
        self.redis_client.setex(f"features:{key}", 3600, json.dumps(features.tolist()))

    def get_cached_features(self, key):
        """从Redis获取缓存的特征数据"""
        import json
        cached = self.redis_client.get(f"features:{key}")
        return json.loads(cached) if cached else None

    def save_model_metrics(self, model_name, metrics):
        """保存模型评估指标到MySQL"""
        with self.get_mysql_connection() as conn:
            cursor = conn.cursor()
            query = """
            INSERT INTO model_metrics (model_name, mae, timestamp)
            VALUES (%s, %s, NOW())
            """
            cursor.execute(query, (model_name, metrics['mae']))
            conn.commit()

    def save_training_data(self, data):
        """保存训练数据到MySQL"""
        with self.get_mysql_connection() as conn:
            cursor = conn.cursor()
            # 创建表（如果不存在）
            create_table_query = """
            CREATE TABLE IF NOT EXISTS training_data (
                id INT AUTO_INCREMENT PRIMARY KEY,
                time DATETIME,
                power_load FLOAT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
            """
            cursor.execute(create_table_query)

            # 插入数据
            insert_query = "INSERT INTO training_data (time, power_load) VALUES (%s, %s)"
            data_to_insert = [(row['time'], row['power_load']) for _, row in data.iterrows()]
            cursor.executemany(insert_query, data_to_insert)
            conn.commit()
